#!/usr/bin/python

import imaplib, email
import sys, getopt
import subprocess
import os
import time
import email.charset

from email.parser import Parser

'''
  a simple script for mail interaction with remote host
  Author: zhaoxiaoqiang007@gmail.com

'''

daemon = False
log_f = sys.stdout
clean = False
address = None
recv_server = None
send_server = None
user = None
passwd = None
dest = None

def usage():
    print "mail_cmd.py [-hdc] <-r receive mail server> [-s send mail server] <-u user> <-p password> [-a address]"

def get_cmd():
    
    
    while True:
        try:
            M = imaplib.IMAP4_SSL(recv_server)
        except Exception,e:
            do_log('Connect error: %s, Retry...'% e)

        else:
            break
    
    try:
        try:
            M.login(user, passwd)
        except Exception,e:
            do_log('login error: %s' %e)
            M.close()
            sys.exit(1)

        do_log("login server \'%s\' success!"% recv_server)
        result, message = M.select("INBOX")
        #print "Find %s message"%message[0]
        result, msgnum = M.search(None, '(UNSEEN (SUBJECT "CMD::"))')
        
        for num in msgnum[0].split():
            result, data = M.fetch(num, '(RFC822)')
            #print "UID: %s"%num
            msg = email.message_from_string(data[0][1])
            subject = msg.get("subject")
            if "CMD::" not in subject:
                continue
            #print "SUBJECT: %s"%msg.get("subject")
            #print "CONTENTS: %s"%msg.get_payload()
            cont = None
            real_msg = None
            if msg.is_multipart():
                for m in msg.get_payload():
                    if m.get_content_type() == "text/plain":
                        real_msg = m
                        break
            else:
                real_msg = msg
            

            cont = real_msg.get_payload(decode=True)

            tmp = msg.get("from")

            l = len(tmp)
            s = tmp.index(' <')
            if ' <' in tmp:
                addr = tmp[s + 2:l - 1]
            else:
                addr = tmp
            #print "FROM: %s"%addr
                
            if clean:
                do_log("Delete msg %s"% num)
                M.store(num, '+FLAGS', '\\Deleted')
                result, msgdel = M.expunge()
            
            return [cont.strip(), addr.strip()]
        return [None, None]
        
    except Exception, e:
        do_log('Imap error: %s' %e)
    M.logout()

    return [None, None]

def exe_cmd(cmds):
    ret = ""
    for cmd in cmds:
        cmd1 = cmd + ' 2>&1'
        proc = subprocess.Popen(cmd1, stdout=subprocess.PIPE, shell=True)
        out, err = proc.communicate()
        head = "Command: %s\n\n"% cmd
        ret = ret + head + out + "\n"
    return ret

def mail_back(cmd, string):

    import smtplib
    from email.mime.text import MIMEText

    global send_server

    host = os.uname()[1]
    string = "On host: %s\n\n%s"% (host, string)
    msg = MIMEText(string,'plain','utf8')
    if len(cmd) > 1:
        msg['Subject'] = "Command output: \"%s\" ..."% cmd[0]
    else:
        msg['Subject'] = "Command output: \"%s\""% cmd[0]

    if address != None:
        src = address
    else:
        src = user

    msg['From'] = src
    msg['To'] = dest

    if send_server == None:
        send_server = "smtp" + recv_server[recv_server.index('.'):]

    s = smtplib.SMTP(send_server)
    s.login(user, passwd)
    s.sendmail(src, dest, msg.as_string())
    s.quit()

def poll_req():

    global dest
    cmd_str, addr = get_cmd()
    dest = addr

    if cmd_str != None:
        
        pid = os.fork()
        if pid == 0:
            do_log("Command:\n%s"% "".join(cmd_str))
            do_log("now execute cmd")
            cmds = cmd_str.splitlines()
            ret = exe_cmd(cmds)
            do_log( "ret:\n%s"%ret)

            do_log("mail result back")
            mail_back(cmds, ret)
            do_log("Done!")
            sys.exit(0)

def do_log(msg):

    log_msg = time.ctime() + ': ' + msg + '\n'
    log_f.write(log_msg)
    log_f.flush()

def do_daemon():
    global log_f
    
    # do the UNIX double-fork magic, see Stevens' "Advanced 
    # Programming in the UNIX Environment" for details (ISBN 0201563177)

    os.umask(0) 

    try: 
        pid = os.fork() 
        if pid > 0:
            # exit first parent
            sys.exit(0) 
    except OSError, e: 
        print >>sys.stderr, "fork #1 failed: %d (%s)" % (e.errno, e.strerror) 
        sys.exit(1)

    # decouple from parent environment
    os.setsid() 


    # do second fork
    try: 
        pid = os.fork() 
        if pid > 0:
            # exit from second parent, print eventual PID before
            print "Daemon PID %d" % pid 
            sys.exit(0) 
    except OSError, e: 
        print >>sys.stderr, "fork #2 failed: %d (%s)" % (e.errno, e.strerror) 
        sys.exit(1)

    os.chdir("/") 

    # open log file
    log_f = open("/var/log/mail_cmd.log", "w")

    do_log("Mail command server started!")

def listen_mail():
    opts, args = getopt.getopt(sys.argv[1:], "cdhr:s:u:p:a:", [])

    global recv_server
    global send_server
    global user
    global passwd
    
    global daemon
    global clean
    global address

    for o, a in opts:
        if o in ("-h"):
            usage()
            exit()
        elif o in ("-d"):
            daemon = True
        elif o in ("-a"):
            address = a
        elif o in ("-c"):
            clean = True
        elif o in ("-r"):
            recv_server = a
        elif o in ("-s"):
            send_server = a
        elif o in ("-u"):
            user = a
        elif o in ("-p"):
            passwd = a

    if recv_server == None or user == None or passwd == None:
        usage()
        exit()

    if daemon:
        do_daemon()
            
    while True:
        poll_req()
        time.sleep(8)



if __name__ == "__main__":
    listen_mail()
    
